# -*- coding: utf-8 -*-
"""
Run simulations and generate data for the figures in the manuscript
"A Systems Approach to Cultural Evolution".

@author: Fredrik Jansson
"""

import numpy as np
import pandas as pd
from cultural_systems import Model, Pool

#%% Measures used in the figures
# Frequency of each trait in the population
def freq(model):
    return sum(model.population.has_traits(),0)

# Number of individuals with each set of traits out of four
def freq_4_traits(model):
    traits = [agent.traits for agent in model.population.agents]
    [t.sort() for t in traits]
    traits = [''.join(str(i) for i in t) for t in traits]
    traits.extend(['','0','1','2','3','01','02','03','12','13','23','012','013','023','123','0123'])
    return (pd.value_counts(traits).sort_index()-1).tolist()

#%% Figure 2: Number of individuals with 0 to 4 traits, over time. Two traits can either be compatible or incompatible.
p = Pool([[np.nan,1,-1,-1],[1,np.nan,-1,-1],[-1,-1,np.nan,1],[-1,-1,1,np.nan]]) # create trait universe
m = Model(p,population_size=100,generation_length=100,p_invent=.001,k=10) # instantiate model
(time,results) = m.repeat(1,200,[freq_4_traits]) # run for 200 time steps, compute measure in every time step
tdf = pd.DataFrame(results,columns=['','0','01','012','0123','013','02','023','03','1','12','123','13','2','23','3']) # number of individuals with each set of traits
tcdf = pd.DataFrame({'0': tdf[''], '1': tdf[['0','1','2','3']].sum(1), '2 compatible': tdf[['01','23']].sum(1), '2 incompatible': tdf[['02','03','12','13']].sum(1), '3': tdf[['012','013','023','123']].sum(1), '4': tdf['0123']}) # calculate an aggregated measure

#%% Figure 5: Mutual support may maintain system configurations over a long time.
results_df = [None] * 3
p = Pool.random(20,[-1,1],[.9,.1]) # 20 traits, out of which 10% are compatible
for n in range(2,5):
    p.set_compatibilities(slice(n),np.ones((n,n))) # 2, 3 or 4 traits are mutually compatible
    m = Model(p,population_size=100,generation_length=100,p_invent=.001,k=10)
    (time,results) = m.repeat(100,200,[freq]) # run for 100*200 time steps, compute measure in every 100 time steps
    results_df[n-2] = pd.DataFrame(np.insert(results[:,0:n],n,np.mean(results[:,n:],1),1),time,columns=['A','B','C','D','other'] if n==4 else ['A','B','C','other'] if n==3 else ['A','B','other'])

#%% Figure 7: Prevalence of cultural systems over time.
p = Pool([[np.nan,1,-1,-1],[1,np.nan,-1,-1],[-1,-1,np.nan,1],[-1,-1,1,np.nan]]) # A and B, are mutually compatible, and so are C and D; all other pairs are mutually incompatible
m = Model(p,population_size=100,generation_length=400,p_invent=.001,k=10)
(time,results) = m.repeat(50,400,[freq_4_traits])
tdf = pd.DataFrame(results,time,columns=['','0','01','012','0123','013','02','023','03','1','12','123','13','2','23','3'])
tgdf = pd.DataFrame({'0': tdf[''], '1: A or B': tdf[['0','1']].sum(1), '1: C or D': tdf[['2','3']].sum(1), '2: A and B': tdf['01'], '2: C and D': tdf['23'], '2: other': tdf[['02','03','12','13']].sum(1), '3': tdf[['012','013','023','123']].sum(1), '4': tdf['0123']})
